# models/series_decomp.py
import paddle
import paddle.nn as nn

class SeriesDecomposition(nn.Layer):
    """滑动平均序列分解 (趋势+季节)"""

    def __init__(self, kernel_size=25):
        super().__init__()
        self.avg_pool = nn.AvgPool1D(
            kernel_size=kernel_size,
            stride=1,
            padding=kernel_size // 2,
            exclusive=False
        )

    def forward(self, x):
        # 输入形状: [batch, seq_len, channels]
        x_t = x.transpose([0, 2, 1])  # [batch, channels, seq_len]
        trend = self.avg_pool(x_t).transpose([0, 2, 1])  # [batch, seq_len, channels]
        seasonal = x - trend
        return trend, seasonal
